import socket
import numpy as np
import time

from lerobot.common.robot_devices.control_configs import ControlPipelineConfig
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
import draccus
import torch


cli_args = ['--robot.type=so101L', '--robot.cameras={}', '--control.type=teleoperate']

cfg = draccus.parse(config_class=ControlPipelineConfig, config_path=None, args=cli_args)
robot = ManipulatorRobot(cfg.robot)


if not robot.is_connected:
    robot.connect()


HOST = '47.109.17.68'
PORT = 8869

# HOST = '0.0.0.0'
# PORT = 65432

HEADER_SIZE = 4
DATA_SHAPE = (6,)

def connect_to_server():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((HOST, PORT))
    s.sendall(b'1')  # 标识为客户端1
    return s

def main():
    sock = connect_to_server()
    print("Sender connected to server")
    
    try:
        # 模拟高频发送 (1000包/秒)
        while True:
            # 生成随机数据
            leader_pos = {}
            for name in robot.leader_arms:
                leader_pos[name] = robot.leader_arms[name].read("Present_Position")
            data = leader_pos['main']
            # data = np.random.randn(*DATA_SHAPE).astype(np.float32)
            
            # 序列化并添加包头
            raw_data = data.tobytes()
            header = len(raw_data).to_bytes(HEADER_SIZE, 'big')
            
            # 高效发送
            sock.sendall(header + raw_data)
            
            # 控制发送频率 (根据需求调整)
            time.sleep(0.001)  # 1000Hz
            
    except KeyboardInterrupt:
        print("\nSender stopped")
    finally:
        sock.close()

if __name__ == "__main__":
    main()